import cv2
from ai_library.components.utils import getMaxRect
import numpy as np
import math
from PIL import Image, ImageDraw, ImageFont

# 颜色阈值
color_dict = {
    "绿色": (np.array([41, 120, 100]), np.array([77, 255, 255])),
    "红色": (np.array([0, 120, 100]), np.array([10, 255, 255])),
    "黄色": (np.array([16, 60, 60]), np.array([40, 255, 255])),
    "橙色": (np.array([11, 120, 100]), np.array([25, 255, 255])),
    "青色": (np.array([78, 120, 100]), np.array([99, 255, 255])),
    "蓝色": (np.array([100, 120, 100]), np.array([124, 255, 255])),
    "紫色": (np.array([125, 120, 100]), np.array([155, 255, 255])),
    "黑色": (np.array([0, 0, 0]), np.array([180, 255, 45]))
}

class FigureRec():
    def __init__(self):
        # 存放图形的识别结果
        self.shape_list = []

    def recFigure(self, img):
        try:
            roi_bgr = getMaxRect(img)   # 获取TFT显示屏区域
            roi = cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2GRAY)
            ret, threshed = cv2.threshold(roi, 120, 240, cv2.THRESH_BINARY)
            cv2.imshow('roi_bgr', threshed)

            for color in color_dict:
                self.colorDivision(color, roi_bgr, color_dict[color][0], color_dict[color][1])

            roi_bgr = self.cv2ImgAddText(roi_bgr, self.shape_list)
            img = self.cv2ImgAddRect(roi_bgr, self.shape_list)
        except:
            print("图形识别出错！！")

        # cv2.imshow('roi_bgr', roi_bgr)
        # cv2.imwrite("result.png", roi_bgr)
        return img

    def calculating_angle(self, p1, p2, p0):
        # 从三个坐标点中计算角度
        # p0 是交点
        x1 = p1[0] - p0[0]
        y1 = p1[1] - p0[1]
        x2 = p2[0] - p0[0]
        y2 = p2[1] - p0[1]
        """
        记各顶点坐标A（x1，y1）、B（x2，y2）、C（x3，y3），以求∠A为例：
    	向量AB=（x2-x1，y2-y1），|AB|=c=√[(x2-x1)²+(y2-y1)²]
    	向量AC=（x3-x1，y3-y1），|AC|=b=√[(x3-x1)²+(y3-y1)²]
    	AB · AC=(x2-x1)(x3-x1)+(y2-y1)(y3-y1)
    	cosA = (AB · AC)/(|AB||AC|) =[(x2-x1)(x3-x1)+(y2-y1)(y3-y1)]/√[(x2-x1)²+(y2-y1)²][(x3-x1)²+(y3-y1)²]
        """
        angle = (x1 * x2 + y1 * y2) / math.sqrt((x1 ** 2 + y1 ** 2) * (x2 ** 2 + y2 ** 2))
        return int(math.acos(angle) * 180 / math.pi)

    def calculating_distance(self, p0, p1):
        # 从已知道的两个点计算两点之间距离
        x1 = p1[0] - p0[0]
        y1 = p1[1] - p0[1]
        dis = int(math.sqrt(x1 ** 2 + y1 ** 2))
        return dis

    def getShape(self, cnt):
        area = cv2.contourArea(cnt)
        # 计算弧长
        arcLength = cv2.arcLength(cnt, True)
        # 以指定的精度近似多边形曲线
        approxCurve = cv2.approxPolyDP(cnt, 0.03 * arcLength, True)
        count = len(approxCurve)
        print("count", count)

        # 三角形判断
        if count == 3:
            # 三个顶点，返回结果是[[507 408]]
            a = approxCurve[0][0]
            b = approxCurve[1][0]
            c = approxCurve[2][0]

            # 三个顶点对应的角度 （单位：度）
            angle_a = self.calculating_angle(b, c, a)
            angle_b = self.calculating_angle(a, c, b)
            angle_c = self.calculating_angle(a, b, c)

            # 最大角
            angle_max = max(angle_a, angle_b, angle_c)
            # 若cosA>0 或 tanA>0（A为最大角），则为锐角三角形，84-96
            # python的cos函数是以弧度作为参数
            if math.cos(math.radians(angle_max)) > 0.05:
                return "锐角三角形"
            elif math.cos(math.radians(angle_max)) < -0.05:
                return "钝角三角形"
            else:
                return "直角三角形"
        # 四边形判断
        elif count == 4:
            # 四个顶点
            a = approxCurve[0][0]
            b = approxCurve[1][0]
            c = approxCurve[2][0]
            d = approxCurve[3][0]

            # 四个顶点对应的角度 （单位：度）
            angle_a = self.calculating_angle(b, d, a)
            angle_b = self.calculating_angle(a, c, b)
            angle_c = self.calculating_angle(b, d, c)
            angle_d = self.calculating_angle(a, c, d)

            # 直线ab边长 （单位：像素）
            ab = self.calculating_distance(a, b)
            # 直线bc边长 （单位：像素）
            bc = self.calculating_distance(b, c)

            # 最大角
            angle_max = max(angle_a, angle_b, angle_c, angle_d)
            if (math.cos(math.radians(angle_max)) > 0.07 or math.cos(math.radians(angle_max)) < -0.07) and abs(
                    ab - bc) < 5:
                return "菱形"
            elif abs(ab - bc) < 5:
                return "正方形"
            else:
                return "长方形"
        # 五角星
        elif count == 10:
            return "五角星"
        else:
            # TODO 弧长和面积的比值
            # 筛选出圆形
            # 圆半径
            r = arcLength / (2 * math.pi)
            pi = area / (r ** 2)
            if abs(pi - math.pi) < 1.:
                # 满足圆的条件
                return "圆形"
            else:
                # n边形
                print('shape', str(count) + "边形", '面积', area)
                return None
                # return str(count) + "边形"

    def cv2ImgAddText(self, img, shapes, textColor=(255, 0, 0)):
        if (isinstance(img, np.ndarray)):  # 判断是否OpenCV图片类型
            img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
        # 创建一个可以在给定图像上绘图的对象
        draw = ImageDraw.Draw(img)
        # 字体的格式
        font = ImageFont.truetype('zh.ttf', 10)
        # 绘制文本
        for shape in shapes:
            color, name, x, y, w, h = shape
            print(color, name, x, y, w, h)
            if shape[1] is not None:
                draw.text((x, y), color + name, fill=textColor, font=font)
        # 转换回OpenCV格式
        return cv2.cvtColor(np.asarray(img), cv2.COLOR_RGB2BGR)

    def cv2ImgAddRect(self, img, shapes):
        for shape in shapes:
            color, name, x, y, w, h = shape
            if shape[1] is not None:
                cv2.rectangle(img, (x, y), (x + w, y + h), (225, 0, 0), 1)
        return img

    def colorDivision(self, name, roi_bgr, lowerb, upperb):
        roi_hsv = cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2HSV)
        mask = cv2.inRange(roi_hsv, lowerb, upperb)
        if lowerb is color_dict["红色"][0]:
            lower_red_2 = np.array([156, 120, 100])
            upper_red_2 = np.array([180, 255, 255])
            mask_1 = cv2.inRange(roi_hsv, lower_red_2, upper_red_2)
            mask = cv2.add(mask, mask_1)
        # cv2.imshow("mask2", mask)
        # 黑白图
        ret, threshed = cv2.threshold(mask, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_TRIANGLE)
        tx, ty, tw, th = cv2.boundingRect(threshed)
        print("区域：", name, tx, ty, tw, th)

        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
        # eroded = cv.erode(threshed,kernel)        #腐蚀图像
        # dilated = cv.dilate(threshed,kernel)      #膨胀图像
        # cv.imshow("Eroded Image",eroded)           #显示腐蚀后的图像
        # cv.imshow("Dilated Image",dilated)         #显示膨胀后的图像
        closed1 = cv2.morphologyEx(threshed, cv2.MORPH_CLOSE, kernel, iterations=1)  # 闭运算1
        # cv2.imshow("closed1 Image", closed1)

        # findContours 查找二进制图像中的轮廓
        contours, hierarchy = cv2.findContours(closed1, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)[-2:]
        print("识别个数为：", len(contours))

        for cnt in contours:
            # 得到子图像
            x, y, w, h = cv2.boundingRect(cnt)

            # 计算面积
            area = cv2.contourArea(cnt)
            width = roi_bgr.shape[1]  # (193, 324, 3)
            height = roi_bgr.shape[0]
            # 过滤点面积比较小的轮廓
            if area < 50 or area > 30000:
                continue
            # 过滤边框（0,0,w,h）
            if x < 10 or y < 10 or (width - 10 <= x <= width) or (height - 10 <= y <= height):
                print("过滤边框", x, y, w, h)
                continue

            shape = self.getShape(cnt)
            if shape is not None:
                print('shape', shape, '面积', area)
                self.shape_list.append([name, shape, x, y, w, h])

        # cv2.imshow('out1', threshed)

if __name__ == "__main__":
    img = cv2.imread('')
    rec_figure = FigureRec()
